########################################################################################

library(data.table)
library(RColorBrewer)
library(ggplot2)
library(argparser)
library(dplyr)

##############################################################################

argp <- arg_parser("Plot the deconstructSigs")
argp <- add_argument(argp, "--input_file" , help="")
argp <- add_argument(argp, "--images_path" , help="")
argp <- add_argument(argp, "--mol_type" , help="")

argv <- parse_args(argp)

input_file <- argv$input_file
images_path <- argv$images_path
mol_type <- argv$mol_type

if(1!=1){

	input_file <- "~/20220915_gastric_multiple/dna_combinePublic/finalPlot/molecular_type/decompose_allUSE_AllMut/combine_SBS96.ratio.addInfo.tsv"
  images_path <- "~/20220915_gastric_multiple/finalPlot/molecular_type/decompose_allUSE_AllMut"
  mol_type <- "GS"

}

##############################################################################

dir.create(images_path , recursive = T)

dat_sig <- data.frame(fread(input_file))

##############################################################################
## 读取突变信号
share_list <- c("all" , "trunk" , "private")
class_order <- c("IM" , "IGC" , "DGC" )
from_list <- c( "NJMU" , "OncoSG" , "TCGA" , "Utokyo" , "TMUCIH" )

##############################################################################
## 只看NJMU的样本
dat_combine_SBS <- subset( dat_sig , From == "NJMU" )
dat_combine_SBS <- subset( dat_combine_SBS , Type != "IM + IGC + DGC" )

##############################################################################
## 分子分型
msi_sample <- unique(subset(dat_combine_SBS , Molecular.subtype=="MSI")$Patient)
gs_sample <- unique(subset(dat_combine_SBS , Molecular.subtype=="GS")$Patient)
cin_sample <- unique(subset(dat_combine_SBS , Molecular.subtype=="CIN")$Patient)

if( mol_type == "MSS" ){
	use_sample <- c(gs_sample , cin_sample)
}else if( mol_type == "MSI" ){
	use_sample <- msi_sample
}else if( mol_type == "GS" ){
	use_sample <- gs_sample
}else if( mol_type == "CIN" ){
	use_sample <- cin_sample
}

dat_combine_SBS <- subset( dat_combine_SBS , Patient %in% use_sample )

## 计算占比
dat_combine_SBS_ratio <- c()
for( share in share_list ){
	for(class in class_order){
		tmp <- subset( dat_combine_SBS , mut_class == share & Class == class )
		
		## 同一患者同一病理类型取中位数
		tmp <- tmp %>%
		group_by( Patient , variable , Class , mut_class ) %>%
		summarize( value = round(median(value)) )

		tmp <- tmp %>%
		group_by( Class , variable , mut_class ) %>%
		summarize( values = sum(value) )
		tmp$Exposures <- tmp$values/sum(tmp$values)
		dat_combine_SBS_ratio <- rbind( dat_combine_SBS_ratio , tmp )
	}
}

############################################
## 突变信号在各个类别占比均小于0.01的信号，归为other
all_sig <- dat_combine_SBS_ratio
all_sig$Sig <- all_sig$variable
all_sig$MutNum <- all_sig$values

tmp_exporsure <- all_sig %>%
group_by(Sig , mut_class , Class) %>%
summarize( Exposures_all = sum(Exposures) )

use_sig <- unique(data.frame(tmp_exporsure[tmp_exporsure$Exposures_all > 0.01,])$Sig)

##############################################################################
## SBS1   SBS5   SBS15  SBS17a SBS17b SBS29  SBS39  SBS3 
## SBS1   SBS2   SBS3   SBS5   SBS13  SBS15  SBS17a SBS17b SBS29
## SBS15:Defective DNA mismatch repair
## SBS17a SBS17b SBS39:unknown
## SBS3:Defective homologous recombination-based DNA damage repair 
## SBS18:Possibly damage by reactive oxygen species.
SBS_Age <- c("SBS1" , "SBS5" , "SBS40")
SBS_apobec <- c("SBS2" , "SBS13")
SBS_repair <- c("SBS15" , "SBS6" , "SBS20" , "SBS26" , "SBS18" , "SBS3" )
SBS_POLE <- c("SBS10a" , "SBS10b")
SBS_Smoke <- c("SBS29")
SBS_gastric <- c("SBS17a" , "SBS17b")
SBS_unkown <- c("SBS39")

sig_order <- c( SBS_Age , SBS_repair , SBS_POLE , SBS_apobec , SBS_Smoke , SBS_gastric , SBS_unkown)
all_sig_use <- subset( all_sig , Sig %in% sig_order )

##############################################################################
## 合并OtherSig
#all_sig <- rbind( all_sig_use , other_sig )
all_sig$Type <- all_sig$mut_class
all_sig$Type <- factor(all_sig$Type , levels = share_list , order = T)
all_sig$Class <- factor(all_sig$Class , levels = class_order , order = T)

############################################
## 颜色
col_age <- c( 
	rgb(red=187,green=109,blue=201,alpha=255,max=255) , 
	rgb(red=140,green=58,blue=154,alpha=255,max=255) , 
	rgb(red=125,green=53,blue=138,alpha=255,max=255) 
	)
col_APOBEC <- c( rgb(red=2,green=100,blue=190,alpha=255,max=255) , rgb(red=2,green=83,blue=155,alpha=255,max=255))
col_repair <- c( 
	rgb(red=151,green=209,blue=148,alpha=255,max=255) , 
	rgb(red=101,green=188,blue=97,alpha=255,max=255) ,
	rgb(red=64,green=142,blue=60,alpha=255,max=255) ,
	rgb(red=231,green=244,blue=230,alpha=255,max=255) ,
	rgb(red=161,green=206,blue=99,alpha=255,max=255) ,
	rgb(red=39,green=88,blue=37,alpha=255,max=255) 
	)
col_smoke <- c( rgb(red=247,green=184,blue=70,alpha=255,max=255) )
col_pole <- c( rgb(red=65,green=62,blue=133,alpha=255,max=255) ,
	rgb(red=48,green=104,blue=141,alpha=255,max=255) 
	)

col_gastric <- c(brewer.pal(6,"Pastel1")[1:2])
col_unkown <- c(brewer.pal(6,"Pastel1")[3])
col_other <- "grey"

col_gastric <- c( "#3f60aa" , "#44c1f0" )
col_APOBEC <- c(brewer.pal(6,"Pastel1")[1:2])



############################################
sig_order_new <- c( 
	paste0( SBS_Age  , " (Age)" ) ,
	paste0( SBS_repair[1:4] , " (Mismatch Repair)" ) ,
	paste0( SBS_repair[5] , " (Oxygen)" ) , 
	paste0( SBS_repair[6] , " (Homologous Repair)" ) , 
	paste0( SBS_POLE , " (Polymerase)" ) , 
	paste0( SBS_gastric , " (Gastric)" ) ,
	paste0( SBS_apobec  , " (Apobec)" ) ,
	paste0( SBS_Smoke , " (Tobacco Chewing)" ) , 
	paste0( SBS_unkown , " (unknown)" ) ,
	"Other"
)

col_sig <- c( col_age , col_repair , col_pole , col_gastric ,  col_APOBEC ,  col_smoke ,col_unkown , col_other )
names(col_sig) <- sig_order_new

############################################
## 改名
all_sig$Sig_New <- as.character( all_sig$Sig )
all_sig$Sig_New[all_sig$Sig %in% SBS_Smoke] <- paste0( all_sig$Sig[all_sig$Sig %in% SBS_Smoke] , " (Tobacco Chewing)" )
all_sig$Sig_New[all_sig$Sig %in% SBS_repair[c(1:4)]] <- paste0( all_sig$Sig[all_sig$Sig %in% SBS_repair[c(1:4)]] , " (Mismatch Repair)" )
all_sig$Sig_New[all_sig$Sig %in% SBS_repair[5]] <- paste0( all_sig$Sig[all_sig$Sig %in% SBS_repair[5]] , " (Oxygen)" )
all_sig$Sig_New[all_sig$Sig %in% SBS_repair[6]] <- paste0( all_sig$Sig[all_sig$Sig %in% SBS_repair[6]] , " (Homologous Repair)" )
all_sig$Sig_New[all_sig$Sig %in% SBS_Age] <- paste0( all_sig$Sig[all_sig$Sig %in% SBS_Age] , " (Age)" )
all_sig$Sig_New[all_sig$Sig %in% SBS_apobec] <- paste0( all_sig$Sig[all_sig$Sig %in% SBS_apobec] , " (Apobec)" )
all_sig$Sig_New[all_sig$Sig %in% SBS_gastric] <- paste0( all_sig$Sig[all_sig$Sig %in% SBS_gastric] , " (Gastric)" )
all_sig$Sig_New[all_sig$Sig %in% SBS_unkown] <- paste0( all_sig$Sig[all_sig$Sig %in% SBS_unkown] , " (unknown)" )
all_sig$Sig_New[all_sig$Sig %in% SBS_POLE] <- paste0( all_sig$Sig[all_sig$Sig %in% SBS_POLE] , " (Polymerase)" )

## Other
all_sig$Sig_New[ !(all_sig$Sig_New %in% sig_order_new) ] <- "Other"

all_sig <- all_sig %>%
group_by( Class , Sig_New , Type  ) %>%
summarize( Exposures = sum(Exposures) )

##############################################################################

if(mol_type == "MSS" | mol_type == "GS" | mol_type =="CIN"){
	other_sig <- c( "SBS2 (Apobec)" , "SBS13 (Apobec)" , "SBS4 (Smoke)" , "SBS39 (unknown)" ,  "SBS20 (Mismatch Repair)" ,  "SBS26 (Mismatch Repair)" )
	other_sig <- c( "SBS4 (Smoke)" , "SBS26 (Mismatch Repair)" , "SBS20 (Mismatch Repair)" , "SBS29 (Tobacco Chewing)")
	all_sig$Sig_New <- ifelse( all_sig$Sig_New %in% other_sig , "Other" , all_sig$Sig_New  )
}else if(mol_type == "MSI"){
	other_sig <- c( "SBS2 (Apobec)" , "SBS13 (Apobec)" , "SBS17a (Gastric)" , "SBS17b (Gastric)" , "SBS29 (Tobacco Chewing)" , "SBS10a (Polymerase)" , "SBS10b (Polymerase)" )
	all_sig$Sig_New <- ifelse( all_sig$Sig_New %in% other_sig , "Other" , all_sig$Sig_New  )
}

all_sig$Sig_New <- factor(all_sig$Sig_New , levels = c(sig_order_new) , order = T)

##############################################################################
## Function

col <- col_sig[ names(col_sig) %in% unique(all_sig$Sig_New) ]

images_name <- paste0(images_path,"/Mutation_Signature.decompose.NJMU.",mol_type,".pdf",sep="")
plot <- ggplot(all_sig,aes(x=Class,y=Exposures,fill=factor(Sig_New))) +
	geom_bar(stat="identity") +
	facet_grid(.~Type) +
	#ggtitle(mol_type) +
	xlab(NULL) +
	theme_bw() +
  theme(
      legend.position = 'right',
      legend.title = element_blank() ,
      panel.grid.major=element_blank(),
      panel.grid.minor=element_blank(),
      panel.background = element_blank(),
      #panel.border = element_blank(),
     	strip.text.x = element_text(size = 15,color="black",face='bold'),
      plot.title = element_text(size = 12,color="black",face='bold'),
      legend.text = element_text(size = 8,color="black",face='bold'),
      axis.text.y = element_text(size = 12,color="black",face='bold'),
      axis.title.x = element_text(size = 12,color="black",face='bold'),
      axis.title.y = element_text(size = 12,color="black",face='bold'),
      axis.text.x = element_text(size = 12,color="black",face='bold') ,
      axis.ticks.length = unit(0.2, "cm") ,
      axis.line = element_line(size = 0.5)) +
	scale_fill_manual(values=c(col))


ggsave(file=images_name,plot=plot,width=8/1.2,height=4/1.2)

images_name <- paste0(images_path,"/Mutation_Signature.decompose.NJMU.",mol_type,".tsv",sep="")
write.table(all_sig , images_name , row.names = F , quote = F , sep = "\t" )

##############################################################################
